import { u8, u16, u32, u64, i64, toU64, toI64, Registers as RegisterEnum, MemoryProtection } from '../types';
import { 
  Elf64_Ehdr, 
  Elf64_Phdr, 
  Elf64_Shdr, 
  Elf64_Sym, 
  ELFMAG0, 
  ELFMAG1, 
  ELFMAG2, 
  ELFMAG3, 
  ELFCLASS64, 
  ELFDATA2LSB, 
  ET_EXEC,
  ET_DYN, 
  EM_X86_64, 
  PT_LOAD, 
  PT_INTERP, 
  PT_GNU_STACK, 
  PF_R, 
  PF_W, 
  PF_X 
} from './types';
import { MemoryManager } from '../memory/MemoryManager';
import { RegisterManager } from '../cpu/Registers';

// ELF加载器类
export class ElfLoader {
  private memoryManager: MemoryManager;
  private registerManager: RegisterManager;
  private loaded: boolean = false;
  private baseAddress: u64 = 0n;
  private entryPoint: u64 = 0n;
  private interpreter: string | null = null;
  private execstack: boolean = true;

  constructor(memoryManager: MemoryManager, registerManager: RegisterManager) {
    this.memoryManager = memoryManager;
    this.registerManager = registerManager;
  }

  /**
   * 验证ELF文件是否支持
   */
  public isSupportedExecutable(fileData: u8[]): boolean {
    // 检查文件大小是否足够包含ELF头部
    if (fileData.length < 64) {
      console.error('ELF file too small');
      return false;
    }

    // 检查ELF魔数
    if (fileData[0] !== ELFMAG0 || 
        fileData[1] !== ELFMAG1 || 
        fileData[2] !== ELFMAG2 || 
        fileData[3] !== ELFMAG3) {
      console.error('Invalid ELF magic number');
      return false;
    }

    // 解析ELF头部
    const ehdr = this.parseElfHeader(fileData);
    if (!ehdr) {
      return false;
    }

    // 检查ELF类型
    if (ehdr.type !== ET_EXEC && ehdr.type !== ET_DYN) {
      console.error('ELF is neither ET_EXEC nor ET_DYN');
      return false;
    }

    // 检查ELF类别（必须是64位）
    if (fileData[4] !== ELFCLASS64) {
      console.error('32-bit ELF not supported');
      return false;
    }

    // 检查字节序
    if (fileData[5] !== ELFDATA2LSB) {
      console.error('Big-endian ELF not supported');
      return false;
    }

    // 检查机器类型（必须是x86_64）
    if (ehdr.machine !== EM_X86_64) {
      console.error('ELF is not AMD64');
      return false;
    }

    return true;
  }

  /**
   * 加载ELF文件
   */
  public loadElf(fileData: u8[], programName: string): boolean {
    try {
      // 解析ELF头部
      const ehdr = this.parseElfHeader(fileData);
      if (!ehdr) {
        return false;
      }

      // 设置基地址和入口点
      this.baseAddress = this.chooseBaseAddress(ehdr, fileData.length);
      this.entryPoint = this.baseAddress + ehdr.entry;

      console.log(`Loading ELF file: ${programName}`);
      console.log(`Base address: 0x${this.baseAddress.toString(16)}`);
      console.log(`Entry point: 0x${this.entryPoint.toString(16)}`);

      // 遍历程序头表
      this.execstack = true;
      let lastEnd = 0n;
      let lastProt = 0;

      for (let i = 0; i < ehdr.phnum; i++) {
        const phdr = this.getProgramHeader(fileData, ehdr, i);
        if (!phdr) {
          continue;
        }

        switch (phdr.type) {
          case PT_GNU_STACK:
            this.execstack = (phdr.flags & PF_X) !== 0;
            break;
          case PT_LOAD:
            lastEnd = this.loadLoadSegment(
              fileData, 
              phdr, 
              lastEnd,
              lastProt
            );
            break;
          case PT_INTERP:
            this.interpreter = this.loadInterpreter(fileData, phdr);
            break;
        }
      }

      // 如果有解释器，加载解释器
      if (this.interpreter) {
        console.log(`Loading interpreter: ${this.interpreter}`);
        // 注意：实际实现中需要加载解释器文件
        // 这里简化处理，实际需要读取解释器文件并加载
      }

      // 设置寄存器状态
      this.setupRegisters();

      this.loaded = true;
      return true;
    } catch (error) {
      console.error('Error loading ELF file:', error);
      return false;
    }
  }

  /**
   * 加载Flat二进制文件
   */
  public loadFlatBinary(fileData: u8[], baseAddress: u64 = BigInt(0x400000)): boolean {
    try {
      // 创建一个模拟的PT_LOAD段
      const phdr = new Elf64_Phdr();
      phdr.type = PT_LOAD;
      phdr.flags = PF_R | PF_W | PF_X;
      phdr.offset = 0n;
      phdr.vaddr = baseAddress;
      phdr.filesz = BigInt(fileData.length);
      phdr.memsz = BigInt(fileData.length) + BigInt(4096); // 额外分配一页

      this.baseAddress = baseAddress;
      this.entryPoint = baseAddress;
      this.execstack = true;

      // 加载段
      this.loadLoadSegment(fileData, phdr, 0n, 0);

      // 设置寄存器
      this.setupRegisters();

      this.loaded = true;
      return true;
    } catch (error) {
      console.error('Error loading flat binary:', error);
      return false;
    }
  }

  /**
   * 获取入口点
   */
  public getEntryPoint(): u64 {
    return this.entryPoint;
  }

  /**
   * 检查是否已加载
   */
  public isLoaded(): boolean {
    return this.loaded;
  }

  /**
   * 解析ELF头部
   */
  private parseElfHeader(fileData: u8[]): Elf64_Ehdr | null {
    try {
      const ehdr = new Elf64_Ehdr();

      // 复制ident数组
      for (let i = 0; i < 16; i++) {
        ehdr.ident[i] = fileData[i];
      }

      // 读取其他字段（小端字节序）
      ehdr.type = this.read16(fileData, 16);
      ehdr.machine = this.read16(fileData, 18);
      ehdr.version = this.read32(fileData, 20);
      ehdr.entry = this.read64(fileData, 24);
      ehdr.phoff = this.read64(fileData, 32);
      ehdr.shoff = this.read64(fileData, 40);
      ehdr.flags = this.read32(fileData, 48);
      ehdr.ehsize = this.read16(fileData, 52);
      ehdr.phentsize = this.read16(fileData, 54);
      ehdr.phnum = this.read16(fileData, 56);
      ehdr.shentsize = this.read16(fileData, 58);
      ehdr.shnum = this.read16(fileData, 60);
      ehdr.shstrndx = this.read16(fileData, 62);

      return ehdr;
    } catch (error) {
      console.error('Failed to parse ELF header:', error);
      return null;
    }
  }

  /**
   * 获取程序头表项
   */
  private getProgramHeader(fileData: u8[], ehdr: Elf64_Ehdr, index: number): Elf64_Phdr | null {
    try {
      const offset = ehdr.phoff + BigInt(index * ehdr.phentsize);
      if (Number(offset) + 56 > fileData.length) { // Elf64_Phdr大小为56字节
        return null;
      }

      const phdr = new Elf64_Phdr();
      phdr.type = this.read32(fileData, Number(offset));
      phdr.flags = this.read32(fileData, Number(offset) + 4);
      phdr.offset = this.read64(fileData, Number(offset) + 8);
      phdr.vaddr = this.read64(fileData, Number(offset) + 16);
      phdr.paddr = this.read64(fileData, Number(offset) + 24);
      phdr.filesz = this.read64(fileData, Number(offset) + 32);
      phdr.memsz = this.read64(fileData, Number(offset) + 40);
      phdr.align = this.read64(fileData, Number(offset) + 48);

      return phdr;
    } catch (error) {
      return null;
    }
  }

  /**
     * 加载段
     */
  private loadLoadSegment(
    fileData: u8[],
    phdr: Elf64_Phdr,
    lastEnd: u64,
    lastProt: u32
  ): u64 {
    const flags = phdr.flags;
    const vaddr = phdr.vaddr + this.baseAddress;
    const memsz = phdr.memsz;
    const offset = phdr.offset;
    const filesz = phdr.filesz;
    const pagesize = 4096;
    const start = this.roundDown(vaddr, pagesize);
    const end = this.roundUp(vaddr + memsz, pagesize);
    const skew = vaddr & BigInt(pagesize - 1);

    console.log(`Loading segment: 0x${vaddr.toString(16)} - 0x${(vaddr + memsz).toString(16)}`);
    console.log(`  Flags: ${flags & PF_R ? 'R' : '.'}${flags & PF_W ? 'W' : '.'}${flags & PF_X ? 'X' : '.'}`);
    console.log(`  File offset: 0x${offset.toString(16)}, size: 0x${filesz.toString(16)}`);

    // 检查段的有效性
    if (offset > fileData.length || filesz > fileData.length || (filesz && offset + filesz > fileData.length)) {
      throw new Error('Invalid segment in ELF file');
    }

    if (end < lastEnd) {
      throw new Error('Program headers are not ordered');
    }

    // 计算内存保护属性
    const prot = 
      (flags & PF_R ? 1 : 0) |  // 读权限
      (flags & PF_W ? 2 : 0) |  // 写权限
      (flags & PF_X ? 4 : 0);   // 执行权限

    // 计算内存保护属性
    let protection: MemoryProtection = 0;
    if (flags & PF_R) protection |= MemoryProtection.READ;
    if (flags & PF_W) protection |= MemoryProtection.WRITE;
    if (flags & PF_X) protection |= MemoryProtection.EXECUTE;

    // 分配内存
    this.memoryManager.map(start, Number(end - start), protection);

    // 复制文件内容到内存
    if (filesz > 0) {
      const segmentData = new Uint8Array(fileData.slice(Number(offset), Number(offset + filesz)));
      this.memoryManager.write(vaddr, segmentData);
    }

    // 如果内存大小大于文件大小，初始化额外的内存为0
    if (memsz > filesz) {
      const zeroData = new Uint8Array(Number(memsz - filesz)).fill(0);
      this.memoryManager.write(vaddr + filesz, zeroData);
    }

    return end;
  }

  /**
   * 加载解释器
   */
  private loadInterpreter(fileData: u8[], phdr: Elf64_Phdr): string {
    const offset = phdr.offset;
    const size = phdr.filesz;
    
    if (offset + size > fileData.length) {
      throw new Error('Invalid interpreter segment');
    }

    // 读取解释器路径（以null结尾的字符串）
    let path = '';
    for (let i = 0; i < size; i++) {
      const byte = fileData[Number(offset) + i];
      if (byte === 0) break;
      path += String.fromCharCode(byte);
    }

    return path;
  }

  /**
   * 选择基地址
   */
  private chooseBaseAddress(ehdr: Elf64_Ehdr, fileSize: number): u64 {
    // 对于动态链接的可执行文件，使用随机基地址
    if (ehdr.type === ET_DYN) {
      // 在实际实现中，这里应该使用ASLR (地址空间布局随机化)
      // 简化实现：使用固定的基地址
      return BigInt(0x400000);
    }
    
    // 对于静态链接的可执行文件，使用0作为基地址
    return 0n;
  }

  /**
   * 设置寄存器
   */
  private setupRegisters(): void {
    // 设置指令指针为入口点
    this.registerManager.setInstructionPointer(this.entryPoint);
    
    // 设置栈指针（假设栈大小为16MB，位于较高地址）
    const stackSize = 16 * 1024 * 1024; // 16MB
    const stackTop = BigInt(0x7ffffff00000);    // 假设的栈顶地址
    const stackBase = stackTop - BigInt(stackSize);
    
    // 分配栈内存
    this.memoryManager.map(stackBase, stackSize, MemoryProtection.READ_WRITE);
    
    // 设置栈指针
    this.registerManager.setStackPointer(stackTop - BigInt(8)); // 减去8以对齐
  }

  /**
   * 读取小端序的16位整数
   */
  private read16(data: u8[], offset: number): u16 {
    return (data[offset] | (data[offset + 1] << 8)) >>> 0;
  }

  /**
   * 读取小端序的32位整数
   */
  private read32(data: u8[], offset: number): u32 {
    return (
      data[offset] | 
      (data[offset + 1] << 8) | 
      (data[offset + 2] << 16) | 
      (data[offset + 3] << 24)
    ) >>> 0;
  }

  /**
   * 读取小端序的64位整数
   */
  private read64(data: u8[], offset: number): u64 {
    const low = this.read32(data, offset);
    const high = this.read32(data, offset + 4);
    return (BigInt(high) << BigInt(32)) | BigInt(low);
  }

  /**
   * 向下取整到页边界
   */
  private roundDown(value: u64, pagesize: number): u64 {
    return value - (value % BigInt(pagesize));
  }

  /**
   * 向上取整到页边界
   */
  private roundUp(value: u64, pagesize: number): u64 {
    return this.roundDown(value + BigInt(pagesize) - BigInt(1), pagesize);
  }
}